from torch import device as torch_device
from torch import cuda as torch_cuda

class BasicConfig(object):

    def __init__(self,
                 data_path,
                 record_path,
                 result_path,
                 use_cuda,
                 do_train,
                 do_valid,
                 use_checkpoint,
                 lr,
                 cpu_num,
                 train_batch_size,
                 test_batch_size,
                 epoch_num,
                 train_max_steps,
                 warm_up_step,
                 log_step,
                 valid_step,
                 max_seq_len,
                 negative_sample_size,
                 checkpoint_step,
                 embedding_dim=768,
                 chosen_topk=3
                 ):
        self.data_path = data_path
        self.record_path = record_path
        self.result_path = result_path

        self.use_cuda = use_cuda
        self.device = torch_device('cpu') if not use_cuda or not torch_cuda.is_available() else torch_device('cuda')
        self.do_train = do_train
        self.do_valid = do_valid
        self.use_checkpoint = use_checkpoint

        self.embedding_dim = embedding_dim

        self.cpu_num = cpu_num
        self.learning_rate = lr
        self.train_batch_size = train_batch_size
        self.test_batch_size = test_batch_size
        self.epoch_num = epoch_num
        self.train_max_steps = train_max_steps
        self.warm_up_step = warm_up_step
        self.log_step = log_step
        self.valid_step = valid_step
        self.checkpoint_step = checkpoint_step
        self.max_seq_len = max_seq_len
        self.negative_sample_size = negative_sample_size

        self.cnn_kernel_size = 2
        self.cnn_out_channels = 32
        self.chosen_topk = chosen_topk